# pip install tabulate: 用于dataframe转markdown
import openai
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from collections import deque
from datetime import datetime, timedelta


# 设置支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']
# 解决负号显示问题
plt.rcParams['axes.unicode_minus'] = False

class LLM:
    '''支持工具调用的大语言模型模块'''
    def __init__(self):
        api_key = '<用户 API Key>'
        base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
        self.model = 'qwen-turbo-latest'
        self.client = openai.OpenAI(
            api_key=api_key,
            base_url=base_url,
        )

    def generate_response(self, messages, tools=None, tool_choice=None):
        '''支持工具调用的生成方法'''
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                tools=tools,
                tool_choice=tool_choice,
                temperature=0.5,
                max_tokens=2000
            )
            return response.choices[0].message
        except Exception as e:
            return f'生成响应时出错：{str(e)}'

class MemoryManager:
    '''增强型记忆管理模块'''
    def __init__(self, max_history=5):
        self.history = deque(maxlen=max_history)
        self.data_cache = {}  # 数据集缓存

    def add_history(self, data):
        '''添加对话历史'''
        self.history.append(data)

    def get_history(self):
        '''获取对话历史'''
        return list(self.history)
    
    def cache_data(self, data_key, dataframe):
        '''缓存数据集'''
        self.data_cache[data_key] = dataframe
    
    def get_data(self, data_key):
        '''获取缓存数据'''
        return self.data_cache.get(data_key, None)

class DataAnalyzerTools:
    '''数据分析工具集'''
    @staticmethod
    def generate_sample_data(data_type='sales'):
        '''生成模拟数据集'''
        np.random.seed(42)
        
        if data_type == 'sales':
            dates = [datetime(2023,1,1) + timedelta(days=i) for i in range(100)]
            return pd.DataFrame({
                '日期': dates,
                '产品': np.random.choice(['A', 'B', 'C'], 100),
                '销售额': np.random.randint(100, 500, 100),
                '销售量': np.random.randint(10, 50, 100),
                '地区': np.random.choice(['华东', '华南', '华北'], 100)
            })
        
        elif data_type == 'user':
            return pd.DataFrame({
                '用户ID': [f'U{1000+i}' for i in range(100)],
                '注册日期': pd.date_range('2025-01-01', periods=100),
                '最后登录': pd.date_range('2025-03-01', periods=100),
                '访问次数': np.random.poisson(5, 100),
                '平均停留(min)': np.random.uniform(1, 15, 100).round(1)
            })
        
        elif data_type == 'inventory':
            return pd.DataFrame({
                '产品ID': [f'P{5000+i}' for i in range(100)],
                '当前库存': np.random.randint(0, 200, 100),
                '安全库存': np.random.randint(50, 100, 100),
                '周销量': np.random.randint(10, 50, 100),
                '补货状态': np.random.choice(['正常', '需补货', '紧急补货'], 100, p=[0.6, 0.3, 0.1])
            })
        return pd.DataFrame()

    @classmethod
    def describe_data(cls, df):
        '''生成统计描述'''
        return df.describe().to_markdown()
    
    @classmethod
    def visualize_data(cls, df, columns):
        '''生成可视化图表'''
        try:
            plt.figure(figsize=(10, 6))
            if len(columns) == 1:
                df[columns[0]].hist()
                plt.title(f'{columns[0]} 分布图')
            else:
                df[columns].plot()
                plt.title('趋势分析')
            
            save_path =f"output_{datetime.now().strftime('%H%M%S')}.png"
            plt.savefig(save_path, format='png', bbox_inches='tight')
            plt.close()
            return f'可视化成功，图表已保存到：{save_path}'
        except Exception as e:
            return f'可视化失败：{str(e)}'

class DataAnalysisAgent:
    '''完整数据分析智能体'''
    def __init__(self):
        self.llm = LLM()
        self.memory = MemoryManager()
        self.tools = [
            {
                'type': 'function',
                'function': {
                    'name': 'load_data',
                    'description': '加载内置数据集（sales=销售数据，user=用户数据，inventory=库存数据）',
                    'parameters': {
                        'type': 'object',
                        'properties': {
                            'data_type': {
                                'type': 'string', 
                                'enum': ['sales', 'user', 'inventory'],
                                'description': '数据集类型'
                            }
                        },
                        'required': ['data_type']
                    }
                }
            },
            {
                'type': 'function',
                'function': {
                    'name': 'describe_data',
                    'description': '生成当前数据集的统计描述',
                    'parameters': {
                        'type': 'object',
                        'properties': {
                            'data_type': {
                                'type': 'string',
                                'description': '数据集标识符'
                            }
                        },
                        'required': ['data_type']
                    }
                }
            },
            {
                'type': 'function',
                'function': {
                    'name': 'visualize_data',
                    'description': '生成指定字段的可视化图表',
                    'parameters': {
                        'type': 'object',
                        'properties': {
                            'data_type': {
                                'type': 'string',
                                'description': '数据集标识符'
                            },
                            'columns': {
                                'type': 'array',
                                'items': {'type': 'string'},
                                'description': '需要可视化的字段列表'
                            }
                        },
                        'required': ['data_type', 'columns']
                    }
                }
            }
        ]
        
        self.system_prompt = {
            'role': 'system',
            'content': '''你是数据分析专家，可以调用以下工具：
1. load_data - 加载内置数据集
2. describe_data - 生成统计描述
3. visualize_data - 生成可视化图表
请遵循以下工作流程：
1. 确认用户需求对应的数据集类型
2. 优先使用缓存数据（可通过历史对话获取）
3. 明确说明分析方法和结论'''
        }

    def _handle_tool_call(self, tool_call):
        '''处理工具调用请求'''
        func_name = tool_call.function.name
        args = eval(tool_call.function.arguments)
        
        try:
            if func_name == 'load_data':
                df = DataAnalyzerTools.generate_sample_data(args['data_type'])
                if not df.empty:
                    self.memory.cache_data(args['data_type'], df)
                    preview = df.head(3)
                    return f"{args['data_type']}数据加载成功（{len(df)}行）\n预览：\n{preview}"
                return '数据加载失败'
            
            elif func_name == 'describe_data':
                df = self.memory.get_data(args['data_type'])
                if df is not None:
                    stats = DataAnalyzerTools.describe_data(df)
                    return f'统计描述：\n{stats}'
                return '请先加载数据集'
            
            elif func_name == 'visualize_data':
                df = self.memory.get_data(args['data_type'])
                if df is not None:
                    res = DataAnalyzerTools.visualize_data(df, args['columns'])
                    return res
                return '数据集不存在'
            
            return '未知工具调用'
        except Exception as e:
            return f'工具执行错误：{str(e)}'

    def process_query(self, query):
        '''处理用户查询的完整流程'''
        # 构建消息历史
        messages = [self.system_prompt] + self.memory.get_history()
        messages.append({'role': 'user', 'content': query})
        self.memory.add_history({'role': 'user', 'content': query})
        # print('messages: ', messages)
        # 第一轮调用：获取工具调用请求
        response = self.llm.generate_response(messages, tools=self.tools)
        
        if hasattr(response, 'tool_calls'):
            # 保存工具调用请求
            self.memory.add_history({
                    'role': 'assistant',
                    'content': '',
                    'tool_calls': [response.tool_calls[0]]
                },)
            
            # 处理工具调用
            tool_response = self._handle_tool_call(response.tool_calls[0])
            self.memory.add_history({
                'role': 'tool',
                'content': tool_response,
                'tool_call_id': response.tool_calls[0].id
            })
            
            # 第二轮调用：生成最终回答
            messages = [self.system_prompt] + self.memory.get_history()
            final_response = self.llm.generate_response(messages)
            # self.memory.add_history('assistant', final_response.content)
            return final_response.content
        
        self.memory.add_history('assistant', response.content)
        return response.content
# 大模型回复测试
def llm_test():
    llm = LLM()
    messages = [{'role': 'user', 'content': '你好'}]
    print(llm.generate_response(messages))

def main():
    '''主运行程序'''
    agent = DataAnalysisAgent()
    
    queries = [
        '请分析最近的销售数据',
        '生成全部销售额数据的统计描述摘要',
        '可视化销售额和销售量的趋势',
        '分析用户行为特征',
        '显示库存补货状态分布'
    ]
    
    for query in queries:
        print(f'[用户] {query}')
        response = agent.process_query(query)
        print(f'[助手] {response}')
        print('-' * 60)

if __name__ == '__main__':
    # 测试大模型
    # llm_test()

    main()